import torch
import torch.nn as nn
import torch.optim as optim
import random
import re

# 准备数据
text = '''亲爱的朋友们，我所站立的地方，就是中国。

　　如果你侧耳倾听，你是否听到了那田间蛙鸣和大漠中的驼铃？

　　俯瞰中国乡村的上空，田间依旧的蛙鸣见证了美丽乡村的嬗变。中华有诗云：“稻花香里说丰年，听取蛙声一片。”从关注“三农”到精准扶贫，时代的巨轮飞速向前时，乡村不曾被抛下。它以非凡的活力实现了自身的蜕变，呈现出一幅山水美、人情美的淳朴画卷。

　　从微观层面上，无数致力于美丽乡村建设的人们成为时代的先锋。耶鲁大学村官秦玥飞带领乡村率先走上信息化的轨道；三十年支教女教师以芳兰振蕙叶之姿，使“支月英”这个名字镌刻在两代人的记忆之中。鲁迅曾说：“无穷的远方，无尽的人们，都与我有关。”正是怀着对乡村大地最深沉的眷恋和责任，无数人与曾经的这方土地有了再难分割的羁绊。从宏观层面上，国家的发展从未遗忘农村，从顶层设计再到基层探索，湘西十八洞村的扶贫经验至今还在传播，三权分置体现的对农民的关怀也仍是人们感念的对象。

　　大风起于青萍之末，知止而后有定，激荡之中而后有静。

　　放眼宏观的国际，新丝路上又响起的阵阵驼铃，传来人类命运共同体的集体乡愁。西方有谚云“独木不成林”，非洲朋友也说“水涨荷花高”，一带一路的“传奇”构想是古老中国全球担当的表征，它不是以中国霸权睥睨一切，而是以开放的怀抱包容世界。

　　当普惠互利共赢成为全球治理的共同呼声，中国成为最勇敢的弄潮儿。当“逆全球化”的言论甚嚣尘上，当地球村的乐观之幕徐徐降下，有人发出“世界怎么了？我们怎么办？”的忧患之问。一带一路成为全球对于中国治理方案的共同期待。中国人的文化自觉与文化自信凝聚成责任担当，古老的丝绸之路是养分之源，全新的发展理念又为其注入了不竭的活力。

　　当蛙声伴着田间的风传遍广阔的美丽乡村，当驼铃载着国际的梦奏响一带一路之歌，你会看见古老中国日益崛起的身姿，它化为一声厚重的龙吟，泅渡了历史的重洋，激荡在七大洲的上空。

　　我所站立的地方，就是我的中国，我是什么，中国便是什么。我怎么样，中国便怎么样。

　　今日少年人，当蛙声伴着驼铃，我们向前迈步的脚步声应当更加铿锵有力才是。

　　世界青年人，也请让我们伸出手，握一握，共同穿越巨变的中国。'''

text = re.sub(r'[^\u4e00-\ufa95，。！？]', '', text)

chars = sorted(list(set(text)))
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}

# 定义 LSTM 模型
class CharLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers):
        super(CharLSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden):
        out, hidden = self.lstm(x, hidden)
        out = self.fc(out[:, -1, :])  # 取序列的最后一个时间步
        return out, hidden

    def init_hidden(self, batch_size):
        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_size)
        c0 = torch.zeros(self.num_layers, batch_size, self.hidden_size)
        return h0, c0


# 模型参数
input_size = len(chars)  # 输入大小（字符种类数）
hidden_size = 128
output_size = len(chars)  # 输出大小（字符种类数）
num_layers = 1

# 重新定义模型结构
model = CharLSTM(input_size, hidden_size, output_size, num_layers)

# 加载模型参数
model.load_state_dict(torch.load("pytorch-lstm-word.pth"))

# 生成文本
def generate_text(model, start_str, length=100):
    hidden = model.init_hidden(batch_size=1)
    chars = list(start_str)

    for p in range(len(chars) - 1):
        x = torch.tensor([[char_to_idx[chars[p]]]]).unsqueeze(2)
        x = torch.zeros(1, 1, input_size).scatter_(2, x, 1)
        _, hidden = model(x, hidden)

    pred = chars[-1]

    for p in range(length):
        x = torch.tensor([[char_to_idx[pred]]]).unsqueeze(2)
        x = torch.zeros(1, 1, input_size).scatter_(2, x, 1)
        output, hidden = model(x, hidden)

        output_dist = output.data.view(-1).exp()
        top_i = torch.multinomial(output_dist, 1)[0]
        pred = idx_to_char[top_i.item()]

        chars.append(pred)

    return ''.join(chars)

# 使用模型进行推理
model.eval()
# 例如：生成文本
print(generate_text(model, start_str="人", length=200))
